"""
Classifier module - Provides various classifier implementations for different scenarios.

This module includes:
1. Base classifiers (CosineClf, LinearClf) - For regular classification tasks
2. Continual classifiers (ContinualClf, MoFClf) - For continual/incremental learning

Factory functions:
- build_classifier: Create any type of classifier
- build_base_classifier: Create a base classifier
- build_continual_classifier: Create a continual classifier
"""

from classifier.base import CosineClf, LinearClf
from classifier.base import build_base_classifier
from classifier.continual import ContinualClf, MoFClf, ABContinualClf
from classifier.ease import CosineLinear
from typing import Any, Dict, Type

# 增量学习分类器
CONTINUAL_CLASSIFIER_REGISTRY: Dict[str, Type[ABContinualClf]] = {
    "continual": ContinualClf,
    "mof": MoFClf,
}


def build_continual_classifier(
    name: str, in_features: int, out_features: int, fc_func: str, **kwargs: Any
) -> ABContinualClf:
    """
    Create a continual learning classifier

    Args:
        name: Classifier type name ('continual', 'mof')
        in_features: Input feature dimension
        out_features: Output feature dimension
        fc_func: Fully connected layer constructor function, used to create classification heads for each task
        **kwargs: Additional keyword arguments, such as:
            - constant_dim: Whether to keep input dimension unchanged
            - freeze_old: Whether to freeze old classification heads
            - init_experts: Initial number of experts (for MoFClf)

    Returns:
        ABContinualClf: Instantiated continual learning classifier
    """
    name = name.lower()
    if name not in CONTINUAL_CLASSIFIER_REGISTRY:
        raise ValueError(
            f"Unknown continual classifier type: {name}. Available types: {list(CONTINUAL_CLASSIFIER_REGISTRY.keys())}"
        )

    cls = CONTINUAL_CLASSIFIER_REGISTRY[name]
    return cls(in_features, out_features, fc_func, **kwargs)


__all__ = [
    # Base classifiers
    "CosineClf",
    "LinearClf",
    # Continual classifiers
    "ContinualClf",
    "MoFClf",
    "ABContinualClf",
    # EASE classifiers
    "CosineLinear",
    # Factory functions
    "build_base_classifier",
    "build_continual_classifier",
]
